import sys
import os
import time
from optparse import OptionParser
from torch.autograd import Variable
import numpy as np
from data.dataset import Data
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
from torch import optim
import cv2
from torchvision import transforms
from eval import eval_net
from unet import UNet
from utils import get_ids, split_ids, split_train_val, get_imgs_and_masks, batch
from visdom import Visdom  # 可视化处理模块
from torch import functional as F

# 可视化app
viz = Visdom()

# 超参数
crop_size = 312  # 切片大小
lr = 0.1
gpu = True,
dir_img = './data/train/'
dir_checkpoint = './checkpoints/'
data_lst = 'train.txt'
save_dir = './result'
model_dir = "checkpoints/CP4.0.pth"
batch_size = 1


def test(model):
    test_img = Data(dir_img, data_lst)
    testloader = torch.utils.data.DataLoader(
        test_img, batch_size=1, shuffle=False)
    nm = np.loadtxt(os.path.join(dir_img, data_lst), dtype=str)
    print(len(testloader), len(nm))
    assert len(testloader) == len(nm)
    if gpu:
        model.cuda()
    model.eval()
    start_time = time.time()
    for i, (data, _) in enumerate(testloader):
        if gpu:
            data = data.cuda()
        data = Variable(data, volatile=True)
        out = model(data)
        # 输出最后层结果
        fuse = out.cpu().data.numpy()[0, 0, :, :]

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        cv2.imwrite(os.path.join(save_dir, '%s' % nm[i].replace("'", '').replace("b", '')), 255-fuse*255)
    print('Overall Time use: ', time.time() - start_time)


if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    net.load_state_dict(torch.load(model_dir))
    test(net)